from typing import Dict, Union, Optional
from threading import Lock

from axiom_boot.conf.manager import Settings
from axiom_boot.di.decorators import service
from axiom_boot.di.dependency import autowired
from axiom_boot.vector.embedding.embedding_interface import EmbeddingInterface
from axiom_boot.vector.embedding.dashscope.dashscope_embedding_config import DashScopeEmbeddingConfig
from axiom_boot.vector.embedding.dashscope.dashscope_service import DashScopeEmbeddingService
from axiom_boot.vector.embedding.huggingface.huggingface_embedding_config import HuggingFaceEmbeddingConfig
from axiom_boot.vector.embedding.huggingface.huggingface_embedding_service import HuggingFaceEmbeddingService

# 定义一个类型别名，用于表示所有支持的配置对象类型
EmbeddingConfig = Union[DashScopeEmbeddingConfig, HuggingFaceEmbeddingConfig]


@service()
class EmbeddingManager:
    """
    一个管理器服务，负责按需创建和缓存 Embedding 实例。
    它现在被重构为从 Settings 中读取预定义的配置档案，并按名称提供服务。
    """

    def __init__(self, settings: Settings = autowired()):
        self._settings = settings
        self._embeddings: Dict[str, EmbeddingInterface] = {}
        self._lock = Lock()

    def get_embedding(self, profile_name: str) -> EmbeddingInterface:
        """
        根据配置档案的名称，获取一个 Embedding 实例。

        Args:
            profile_name: 在配置中定义的模型档案名称 (例如, "default")。

        Returns:
            一个符合配置的 EmbeddingInterface 实例。
        """
        if profile_name not in self._embeddings:
            with self._lock:
                if profile_name not in self._embeddings:
                    self._embeddings[profile_name] = self._create_embedding(profile_name)
        
        return self._embeddings[profile_name]

    def get_embedding_dim(self, profile_name: str) -> int:
        """
        根据配置档案名称，获取其向量维度。
        """
        profile_config = self._get_profile_config(profile_name)
        return profile_config.dim

    def _create_embedding(self, profile_name: str) -> EmbeddingInterface:
        """
        根据配置档案，创建一个新的 Embedding 服务实例。
        """
        profile_config = self._get_profile_config(profile_name)
        provider = profile_config.provider

        if provider == "dashscope":
            config = DashScopeEmbeddingConfig(
                model_name=profile_config.model_name,
                api_key=profile_config.api_key
            )
            return DashScopeEmbeddingService(config)
        
        if provider == "huggingface":
            config = HuggingFaceEmbeddingConfig(
                model_name=profile_config.model_name,
                # device 参数是可选的，因此我们安全地 .get()
                device=profile_config.dict().get("device")
            )
            return HuggingFaceEmbeddingService(config)
        
        raise TypeError(f"不支持的 Embedding provider: {provider}")

    def _get_profile_config(self, profile_name: str) -> Dict:
        """从 Settings 对象中安全地获取指定名称的配置档案。"""
        try:
            profile_config = self._settings.embedding.profiles.get(profile_name)
            if not profile_config:
                raise KeyError
            return profile_config
        except (AttributeError, KeyError):
            raise ValueError(f"未在配置中找到名为 '{profile_name}' 的 Embedding 模型档案。") 